{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {},
      "outputs": [],
      "source": [
        "# 2x 2 chips (4 cores) per process:\n",
        "os.environ[\"TPU_CHIPS_PER_HOST_BOUNDS\"] = \"1,2,1\"\n",
        "os.environ[\"TPU_HOST_BOUNDS\"] = \"1,1,1\"\n",
        "# Different per process:\n",
        "os.environ[\"TPU_VISIBLE_DEVICES\"] = \"0,1\" # Change to \"2,3\" for the second machine\n",
        "# Pick a unique port per process\n",
        "os.environ[\"TPU_MESH_CONTROLLER_ADDRESS\"] = \"localhost:8476\"\n",
        "os.environ[\"TPU_MESH_CONTROLLER_PORT\"] = \"8476\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "[TpuDevice(id=0, process_index=0, coords=(0,0,0), core_on_chip=0), TpuDevice(id=1, process_index=0, coords=(0,0,0), core_on_chip=1), TpuDevice(id=2, process_index=0, coords=(0,1,0), core_on_chip=0), TpuDevice(id=3, process_index=0, coords=(0,1,0), core_on_chip=1)]\n"
          ]
        }
      ],
      "source": [
        "import jax\n",
        "print(jax.devices())"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "include_colab_link": true,
      "name": "scratchpad",
      "provenance": []
    },
    "interpreter": {
      "hash": "31f2aee4e71d21fbe5cf8b01ff0e069b9275f58929596ceb00d14d90e3e16cd6"
    },
    "kernelspec": {
      "display_name": "Python 3.8.10 64-bit",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.8.10"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 2
}
